#include <iostream>
#include <cmath>
#include <algorithm>
#include <Eigen/Dense>
#include "npdm_optimal_irreps.h"
#include "npdm_typedefs.h"

namespace Npdm {
//----------------------------------------------------------------------------------------------------------------------------------------------------------
// Build all distinct pairs of creation and destruction operator strings

void NpdmBlueprint::buildIrrepOpStrings( std::vector< std::vector< std::tuple<int,CRE_DES> > > & lhsOpStrings,
                                         std::vector< std::vector< std::tuple<int,CRE_DES> > > & rhsOpStrings ) {

  pdmOrder_ = lhsOpStrings[0].size();

  // Build set of irreducible permutations with a triangular loop
  for (int p=0; p < lhsOpStrings.size(); p++) {
    std::vector< std::tuple<int,CRE_DES> > lhs(lhsOpStrings[p]);

    for (int q=p; q < rhsOpStrings.size(); q++) {
      std::vector< std::tuple<int,CRE_DES> > rhs(rhsOpStrings[q]);

      // Store all-index operator
      std::vector< std::tuple<int,CRE_DES> > allIndex;
      allIndex.reserve( lhs.size() + rhs.size() );
      allIndex.insert( allIndex.end(), lhs.begin(), lhs.end() );
      allIndex.insert( allIndex.end(), rhs.begin(), rhs.end() );
      idxSet_.insert( allIndex );

    };
  };

  std::cout << "PDM order  = " << pdmOrder_ << std::endl;
  std::cout << "N sites    = " << kmax_ << std::endl;
  std::cout << "Irreps     = " << idxSet_.size() << std::endl;
  std::cout << "(2*k)**P   = " << std::pow(2*kmax_,pdmOrder_) << std::endl;
  std::cout << "Ratio      = " << int( (idxSet_.size()*100) / std::pow(2*kmax_,pdmOrder_) ) << "%" << std::endl;

//  // Analytical formula for number of tuples
////  int N = kmax_;
//// Formula for spatial orbitals
////  long int tmp = N*(N+1)*(N+2)*(N+3)*(N*N +2)/72;
//  std::cout << "Theory  = " << tmp << std::endl;
//  assert (tmp == ijklmn_.size());

}

//----------------------------------------------------------------------------------------------------------------------------------------------------------
// Permute input operator string to canonical "DMRG form" (defined by std::sort) and return parity of permutation matrix

int NpdmBlueprint::getPermParity( const std::vector< std::tuple<int,CRE_DES> > & opString) { 

////FIXME parity thing must also know if we're preparing transposed DMRG form or not
  return 0;
} 

//  std::vector< std::pair< std::tuple<int,CRE_DES,ALPHA_BETA>, std::vector<int> > > opMatString;
//
//
//  // Initialize permutation matrix for opString
//  std::vector< std::tuple<int,CRE_DES,ALPHA_BETA> >::const_iterator op;
//  int i = 0;
//  for (op = opString.begin(); op != opString.end(); op++) {
//    std::vector<int> idx(opString.size(),0);
//    idx[i] = 1;
//    opMatString.push_back( std::make_pair( *op, idx) );
//    i++;
//  };
//
//  // Permute to build permutation matrix
//  std::sort( opMatString.begin(), opMatString.end() );
//
//  // Back out permutation matrix as explicit matrix type
//  Eigen::MatrixXi permMat(opString.size(),opString.size());
//  for (int i=0; i<opMatString.size(); i++) {
//    for (int j=0; j<opMatString.size(); j++) {
//      permMat(i,j) = opMatString[i].second[j];
//    };
//  };
//
//  // Parity is just the determinant of permutation matrix
//  int det = permMat.determinant();
//  assert(abs(det) == 1);
//
//  return det;
//}

//----------------------------------------------------------------------------------------------------------------------------------------------------------
// DMRG ordered form is unique up to a cre/dest permutation (transpose) of the original irreducible string. 
// Should we use the transpose or not for each string considered?
//
// Would like to use the combinations that minimize the number of (+,.+,...) patterns?
// How to do this?
//

void NpdmBlueprint::sortIntoDMRGform( std::vector< std::tuple<int,CRE_DES> > & dmrgString) {

std::sort( dmrgString.begin(), dmrgString.end() );
return;
}

/*
  // Sort original string into DMRG form and evaluate positions of creation operators
  std::vector< std::tuple<int,CRE_DES> > opString, trOpString;
  opString = dmrgString;
  std::sort( opString.begin(), opString.end() );
  int eval = 0;
  for (int i=0; i < opString.size(); i++) {
    if (std::get<1>(opString[i]) == CREATION) eval += i*i;
  };

  // Tranpose input string
  int idx;
  ALPHA_BETA ab;
  CRE_DES cre_des;
  trOpString.clear();
  for (int i=dmrgString.size(); i-- > 0; ) {
    std::tie(idx,cre_des,ab) = dmrgString[i];
    if (cre_des == CREATION) cre_des = DESTRUCTION;
    else cre_des = CREATION;
    trOpString.push_back( std::make_tuple(idx,cre_des,ab) );
  }; 
  
  // Sort transposed string into DMRG form and evaluate positions of creation operators
  std::sort( trOpString.begin(), trOpString.end() );
  int trEval = 0;
  for (int i=0; i < trOpString.size(); i++) {
    if (std::get<1>(trOpString[i]) == CREATION) trEval += i*i;
  };

  // Use the sorted form (with transpose or not) which has most creation ops to left
  if ( trEval < eval ) dmrgString = trOpString;
  else std::sort( dmrgString.begin(), dmrgString.end() );

}
*/

//----------------------------------------------------------------------------------------------------------------------------------------------------------
// Sort each operator string by site into DMRG form and store with parity sign factor
// Note that creation ops are sorted to the left wrt destruction, so these permutations can only produce a sign change at most

void NpdmBlueprint::buildDMRGmap( ) {
 
  std::vector< std::tuple<int,CRE_DES> > dmrgString;
  std::set< std::vector< std::tuple<int,CRE_DES> > >::const_iterator opString;

  idxMap_.clear();
  for (opString = idxSet_.begin(); opString != idxSet_.end(); opString++) {
    dmrgString = *opString;
    // Sort by site index into DMRG form
    sortIntoDMRGform( dmrgString );
    // Get sign change on permuting from original string to DMRG form
    int parity = getPermParity( *opString );
    idxMap_[ dmrgString ] = std::make_pair( *opString, parity );
  };

  // Bin DMRG-ordered operators according to position in DMRG sweep they should be calculated.
  dmrgIdxBins_.clear();
  std::map< std::vector< std::tuple<int,CRE_DES> > , 
            std::pair< std::vector< std::tuple<int,CRE_DES> >, int > >::const_iterator mapIter;
  for (mapIter = idxMap_.begin(); mapIter != idxMap_.end(); mapIter++) {
    // NOTE!
    // Specific site index of operator in the DMRG-ordered operator string used as the multimap key
    int dotIdx = std::max(2, std::get<0>( (mapIter->first)[pdmOrder_] ) );
    dmrgIdxBins_.insert( std::make_pair( dotIdx , mapIter->first ) );
  };

}

//----------------------------------------------------------------------------------------------------------------------------------------------------------

std::set< std::vector< std::tuple<int,CRE_DES> > > NpdmBlueprint::getSpaceStringsForSite(int dotSite) {

  std::pair< std::multimap< int, std::vector< std::tuple<int,CRE_DES> > >::iterator,
             std::multimap< int, std::vector< std::tuple<int,CRE_DES> > >::iterator > range;
  std::multimap< int, std::vector< std::tuple<int,CRE_DES> > >::iterator iter;
  std::set< std::vector< std::tuple<int,CRE_DES> > > spaceStrings;

  range = dmrgIdxBins_.equal_range(dotSite);
  for ( iter = range.first; iter != range.second; iter++ ){
    spaceStrings.insert( iter->second );
  };

  return spaceStrings;
}

//----------------------------------------------------------------------------------------------------------------------------------------------------------
// Get block operators for LHS, Dot and RHS at given (dot)site

std::map< char, std::set< std::vector< std::tuple<int,CRE_DES> > > > NpdmBlueprint::getBlockOpsForSite(int dotSite) {

  std::pair< std::multimap< int, std::vector< std::tuple<int,CRE_DES> > >::const_iterator,
             std::multimap< int, std::vector< std::tuple<int,CRE_DES> > >::const_iterator > range;
  std::multimap< int, std::vector< std::tuple<int,CRE_DES> > >::const_iterator iter;
  std::map< char, std::set< std::vector< std::tuple<int,CRE_DES> > > > blockSet;
  std::map< char, std::vector< std::tuple<int,CRE_DES> > > blockMap;

  range = dmrgIdxBins_.equal_range(dotSite);
  for ( iter = range.first; iter != range.second; iter++ ){
    blockMap = splitIntoBlocks(iter->first, iter->second);
    blockSet['l'].insert( blockMap['l'] );
    blockSet['d'].insert( blockMap['d'] );
    blockSet['r'].insert( blockMap['r'] );
  };

  return blockSet;
}

//----------------------------------------------------------------------------------------------------------------------------------------------------------
// Split DMRG-ordered operator string into Left, Dot and Right pieces

std::map< char, std::vector< std::tuple<int,CRE_DES> > > 
NpdmBlueprint::splitIntoBlocks( int dotSite, const std::vector< std::tuple<int,CRE_DES> > & opString) {

  std::map< char, std::vector< std::tuple<int,CRE_DES> > > blockMap;
  char ldr[3] = {'l','d','r'};

  std::vector< std::tuple<int,CRE_DES> >::const_iterator op;
  
  int key;
  for (op = opString.begin(); op != opString.end(); op++) {

    // Assign each operator to either LHS, RHS or Dot based on its site index compared to dotSite
    key = std::max( std::min( std::get<0>(*op) - dotSite +1 ,2), 0); 
    blockMap[ ldr[key] ].push_back(*op);

    // Asserts
    assert(key >= 0); assert(key <= 2);
    if (dotSite > 2) assert( blockMap['l'].size() < pdmOrder_+1 );
    assert( blockMap['r'].size() < pdmOrder_ );
  };
  
  return blockMap;
}

//----------------------------------------------------------------------------------------------------------------------------------------------------------

void NpdmBlueprint::printIdx( const std::vector< std::tuple<int,CRE_DES> > & opString) {

  std::vector< std::tuple<int,CRE_DES> >::const_iterator op;
  char ab, cre_des;

  for (op = opString.begin(); op != opString.end(); op++) {
    if (std::get<1>(*op) == CREATION) cre_des = '+';
    else cre_des = '.';
    std::cout << "(" << std::get<0>(*op) << "," << cre_des << ") " ;
  };
  std::cout << std::endl;

}

//----------------------------------------------------------------------------------------------------------------------------------------------------------

void NpdmBlueprint::printIrreps() {

  std::cout << "\nIrreducible permutations:\n";
  std::cout << "---------------\n";

  std::set< std::vector< std::tuple<int,CRE_DES> > >::const_iterator opString;
  for (opString = idxSet_.begin(); opString != idxSet_.end(); opString++) {
    printIdx( *opString );
  };
  std::cout << "---------------\n\n";

}

//----------------------------------------------------------------------------------------------------------------------------------------------------------

void NpdmBlueprint::printDMRGmap( ) {

  // PRINT OUT BY DOT-SITE INDEX
  std::pair< std::multimap< int, std::vector< std::tuple<int,CRE_DES> > >::const_iterator,
             std::multimap< int, std::vector< std::tuple<int,CRE_DES> > >::const_iterator > range;

  std::multimap< int, std::vector< std::tuple<int,CRE_DES> > >::const_iterator iter;
  std::map< char, std::vector< std::tuple<int,CRE_DES> > > blockMap;

  for (int k=2; k<=kmax_; k++) {
    std::cout << "--------------------------------------\n";
    std::cout << "Dot site = " << k << std::endl;
    range = dmrgIdxBins_.equal_range(k);
    for ( iter = range.first; iter != range.second; iter++ ){
      printIdx(iter->second);
      blockMap = splitIntoBlocks(iter->first, iter->second);
//      if ( blockMap['d'].size() != 3 ) continue;
//      if ( blockMap['l'].size() > 0 ) continue;
      // Print out LHS, DOT, RHS blocks separately
      std::cout << "LHS = ";
      printIdx(blockMap['l']);
      std::cout << "Dot = ";
      printIdx(blockMap['d']);
      std::cout << "RHS = ";
      printIdx(blockMap['r']);
      std::cout << std::endl;
    };
  };
}

//----------------------------------------------------------------------------------------------------------------------------------------------------------

void NpdmBlueprint::printBlockOps() {

  std::map< char, std::set< std::vector< std::tuple<int,CRE_DES> > > > blockSet;

  for (int i=2; i<= kmax_; i++) {
    blockSet = getBlockOpsForSite(i);
    std::set< std::vector< std::tuple<int,CRE_DES> > >::const_iterator setIter;

    std::cout << "------------------------------------\n";
    std::cout << "Site = " << i << std::endl;
    std::cout << "------------------------------------\n";
    std::cout << "LHS ops:\n";
    for ( setIter = blockSet['l'].begin(); setIter != blockSet['l'].end(); setIter++ ){
      if (setIter->size() > 0) printIdx(*setIter);
    };
    std::cout << "Dot ops:\n";
    for ( setIter = blockSet['d'].begin(); setIter != blockSet['d'].end(); setIter++ ){
      if (setIter->size() > 0) printIdx(*setIter);
    };
    std::cout << "RHS ops:\n";
    for ( setIter = blockSet['r'].begin(); setIter != blockSet['r'].end(); setIter++ ){
      if (setIter->size() > 0) printIdx(*setIter);
    };
  };

}

//----------------------------------------------------------------------------------------------------------------------------------------------------------
}
